%% HJB SOLVER FOR THE OPTIMAL CONTRACT
function [V,lcp_flag] = optimal_hjb_solver(V00,C, Beta, Eta, K, f, ExRet, mu_v, mu_y, vol_y, vol_phi, iter, min_iter_lcp, params, grid_params, num_params)

%% PARAMTERS
%parameters
temp = num2cell(params);
[rho, r, ra, alpha, sigma, lambda, M, eta_min, eta_max] = temp{:};

[n_grid, Y_mat, PHI_mat, dy, dphi] = grid_params{:};

[max_iter, D, thr] = num_params{:};

%% v-term
mat_v = reshape(mu_v, n_grid*n_grid, 1);

AA_v = spdiags(mat_v, 0, n_grid*n_grid,n_grid*n_grid);

%% FIRST DERIVATIVES
up_mu_y = max(mu_y, 0)/dy;
low_mu_y = -min(mu_y, 0)/dy;
cent_mu_y = -up_mu_y - low_mu_y;

up_mat_y=0; %This is needed because of the peculiarity of spdiags.
for j=1:n_grid
    up_mat_y=[up_mat_y;up_mu_y(1:n_grid-1,j);0];
end

cent_mat_y=reshape(cent_mu_y,n_grid*n_grid,1);


low_mat_y=low_mu_y(2:n_grid,1);
for j=2:n_grid
    low_mat_y=[low_mat_y;0;low_mu_y(2:n_grid,j)];
end

%% AA matrix for v-term and first derivatives
AA_y =spdiags(cent_mat_y,0,n_grid*n_grid,n_grid*n_grid)+ ...
    spdiags(up_mat_y,1,n_grid*n_grid,n_grid*n_grid)+...
    spdiags(low_mat_y,-1,n_grid*n_grid,n_grid*n_grid);

AA = AA_v + AA_y;

%% Second derivatives
vol2_y = vol_y.^2/(2*dy^2); %divided by 2!!!
vol2_phi = vol_phi.^2/(2*dphi^2);

up_mat_vol_y=0; %This is needed because of the peculiarity of spdiags.
for j=1:n_grid
    up_mat_vol_y=[up_mat_vol_y;vol2_y(1:n_grid-1,j);0];
end

cent_mat_vol_y=reshape(-vol2_y*2,n_grid*n_grid,1); %Mutiply by 2!!!

low_mat_vol_y=vol2_y(2:n_grid,1);
for j=2:n_grid
    low_mat_vol_y=[low_mat_vol_y;0;vol2_y(2:n_grid,j)];
end

up_mat_vol_phi=zeros(n_grid,1); %This is necessary because of the peculiar way spdiags is defined.
for j=1:n_grid
    up_mat_vol_phi=[up_mat_vol_phi;vol2_phi(1:n_grid,j)];
end

cent_mat_vol_phi = reshape(-vol2_phi*2, n_grid*n_grid, 1);

low_mat_vol_phi=vol2_phi(1:n_grid,2);
for j=3:n_grid
    low_mat_vol_phi=[low_mat_vol_phi;vol2_phi(1:n_grid,j)];
end

BB_yy =spdiags(cent_mat_vol_y,0,n_grid*n_grid,n_grid*n_grid)+ ...
    spdiags(up_mat_vol_y,1,n_grid*n_grid,n_grid*n_grid)+...
    spdiags(low_mat_vol_y,-1,n_grid*n_grid,n_grid*n_grid);

BB_phiphi =spdiags(cent_mat_vol_phi,0,n_grid*n_grid,n_grid*n_grid)+...
    spdiags(low_mat_vol_phi,-n_grid,n_grid*n_grid,n_grid*n_grid)+...
    spdiags(up_mat_vol_phi,n_grid,n_grid*n_grid,n_grid*n_grid);

BB = BB_yy + BB_phiphi;

%% "Cross derivatives" with Scaling Variable
cross_y = Beta.*vol_y/(2*dy);
up_cross_y = cross_y;
low_cross_y = -cross_y;

cross_phi = vol_phi.*Beta/(2*dphi);
up_cross_phi = cross_phi;
low_cross_phi = -cross_phi;

up_mat_cross_y=0; %This is needed because of the peculiarity of spdiags.
for j=1:n_grid
    up_mat_cross_y=[up_mat_cross_y;up_cross_y(1:n_grid-1,j);0];
end

low_mat_cross_y=low_cross_y(2:n_grid,1);
for j=2:n_grid
    low_mat_cross_y=[low_mat_cross_y;0;low_cross_y(2:n_grid,j)];
end

up_mat_cross_phi=zeros(n_grid,1); %This is necessary because of the peculiar way spdiags is defined.
for j=1:n_grid
    up_mat_cross_phi=[up_mat_cross_phi;up_cross_phi(1:n_grid,j)];
end

low_mat_cross_phi=low_cross_phi(1:n_grid,2);
for j=3:n_grid
    low_mat_cross_phi=[low_mat_cross_phi;low_cross_phi(1:n_grid,j)];
end

CC_phi = spdiags(up_mat_cross_phi,n_grid,n_grid*n_grid,n_grid*n_grid)+...
    + spdiags(low_mat_cross_phi,-n_grid,n_grid*n_grid,n_grid*n_grid);    

CC_y = spdiags(up_mat_cross_y,1,n_grid*n_grid,n_grid*n_grid)+...
        + spdiags(low_mat_cross_y,-1,n_grid*n_grid,n_grid*n_grid);    

CC = CC_phi + CC_y;

%% Cross Derivatives
cross = vol_y.*vol_phi/(dy*dphi*2); % already divided by two!
pos_cross = max(cross,0);
neg_cross = min(cross,0);

cent_cent_cross = 2*(pos_cross - neg_cross);
up_up_cross = pos_cross;
low_low_cross = pos_cross;
up_cent_cross = -pos_cross + neg_cross;
low_cent_cross = -pos_cross + neg_cross;
cent_up_cross = -pos_cross + neg_cross;
cent_low_cross = -pos_cross + neg_cross;
up_low_cross = -neg_cross;
low_up_cross = -neg_cross;

cent_cent_mat = reshape(cent_cent_cross, n_grid*n_grid,1);

up_cent_mat=0; %This is needed because of the peculiarity of spdiags.
for j=1:n_grid
    up_cent_mat=[up_cent_mat;up_cent_cross(1:n_grid-1,j);0];
end

low_cent_mat=low_cent_cross(2:n_grid,1);
for j=2:n_grid
    low_cent_mat=[low_cent_mat;0;low_cent_cross(2:n_grid,j)];
end

cent_up_mat=zeros(n_grid,1); %This is necessary because of the peculiar way spdiags is defined.
for j=1:n_grid
    cent_up_mat=[cent_up_mat;cent_up_cross(1:n_grid,j)];
end

cent_low_mat=cent_low_cross(1:n_grid,2);
for j=3:n_grid
    cent_low_mat=[cent_low_mat;cent_low_cross(1:n_grid,j)];
end

up_up_mat = [0;zeros(n_grid,1)];
for j=1:n_grid
    up_up_mat=[up_up_mat;up_up_cross(1:n_grid-1,j);0];
end    

low_low_mat = low_low_cross(2:end,2);
for j=3:n_grid
    low_low_mat=[low_low_mat;0;low_low_cross(2:n_grid,j)];
end

up_low_mat = [0;up_low_cross(1:n_grid-1,2);0];
for j=3:n_grid
    up_low_mat=[up_low_mat;up_low_cross(1:n_grid-1,j);0];
end

low_up_mat = [zeros(n_grid,1);low_up_cross(2:n_grid,1);];
for j=2:n_grid
    low_up_mat=[low_up_mat;0;low_up_cross(2:n_grid,j)];
end

DD = spdiags(cent_cent_mat,0,n_grid*n_grid,n_grid*n_grid) +...
        spdiags(up_cent_mat,1,n_grid*n_grid,n_grid*n_grid)+...
        spdiags(low_cent_mat,-1,n_grid*n_grid,n_grid*n_grid)+...
        spdiags(cent_up_mat,n_grid,n_grid*n_grid,n_grid*n_grid)+...
        spdiags(cent_low_mat,-n_grid,n_grid*n_grid,n_grid*n_grid)+...
        spdiags(up_up_mat,n_grid+1, n_grid*n_grid,n_grid*n_grid)+...
        spdiags(low_low_mat,-n_grid-1, n_grid*n_grid,n_grid*n_grid)+...
        spdiags(up_low_mat, -n_grid+1, n_grid*n_grid,n_grid*n_grid)+...
        spdiags(low_up_mat,n_grid-1, n_grid*n_grid,n_grid*n_grid);


%% Adjustment at the phi boundary
phi_adj_mat = zeros(n_grid,n_grid);
phi_adj_mat(:,1) = low_cross_phi(:,1)  + vol2_phi(:,1) + up_low_cross(:,1)+...
                    cent_low_cross(:,1) + low_low_cross(:,1);
phi_adj_mat(:,end) = up_cross_phi(:,end)  + vol2_phi(:,end) + up_up_cross(:,end) +...
                    cent_up_cross(:,end) + low_up_cross(:,end);

phi_adj_vect = reshape(phi_adj_mat, n_grid*n_grid,1);

phi_Adj = spdiags(phi_adj_vect,0, n_grid*n_grid, n_grid*n_grid);


%% Adjustment at the y boundary
y_adj_mat = zeros(n_grid, n_grid);

y_adj_mat(end,:) = up_mu_y(end,:) + vol2_y(end,:) + up_cross_y(end,:) +...
                   up_low_cross(end,:) + up_cent_cross(end,:) + up_up_cross(end,:); 

y_adj_vect = reshape(y_adj_mat, n_grid*n_grid,1);

y_Adj = spdiags(y_adj_vect,0, n_grid*n_grid, n_grid*n_grid);

%% Total adjustment
Adj = phi_Adj+ y_Adj;

%% New value function
P = AA + BB + CC + DD + Adj;

f_stacked = reshape(f,n_grid*n_grid,1);
Q = ((r + 1/D)*speye(n_grid*n_grid) - P);

V00_stacked = reshape(V00,n_grid*n_grid,1);

V_stacked = Q\(f_stacked + V00_stacked/D);
V = reshape(V_stacked,n_grid, n_grid);
